package com.coocaa.typhoon.api.third.volcengine.service;

import cn.hutool.crypto.SecureUtil;
import com.alibaba.fastjson.JSON;
import com.coocaa.cirrus.common.base.util.JsonUtil;
import com.coocaa.cirrus.common.base.util.RetryUtil;
import com.coocaa.cirrus.common.third.filestore.config.TosProperties;
import com.coocaa.typhoon.api.third.volcengine.model.VikingDbUserFileEntity;
import com.coocaa.typhoon.api.third.volcengine.model.output.UserFileVO;
import com.coocaa.typhoon.api.third.volcengine.util.VikingDbConvert;
import com.coocaa.typhoon.api.third.volcengine.yml.VolcConfigYml;
import com.coocaa.typhoon.api.v1.userfile.model.input.GetSearchDataModel;
import com.volcengine.service.vikingDB.CollectionClient;
import com.volcengine.service.vikingDB.IndexClient;
import com.volcengine.service.vikingDB.VikingDBService;
import com.volcengine.service.vikingDB.common.*;
import lombok.extern.slf4j.Slf4j;
import org.apache.commons.collections4.CollectionUtils;
import org.springframework.beans.factory.InitializingBean;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.stereotype.Service;

import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;

/**
 * @author bijiahao
 * @date : 2025/9/18 14:21
 * @description
 */
@Slf4j
@Service
public class VikingDbUserFileService implements InitializingBean {
    @Autowired
    private VikingDBService vikingDBService;
    @Autowired
    private VolcConfigYml volcConfigYml;
    private CollectionClient collectionClient;
    private IndexClient indexClient;
    @Autowired
    private TosProperties tosProperties;

    @Override
    public void afterPropertiesSet() throws Exception {
        collectionClient = new CollectionClient(volcConfigYml.getCollectionName(), volcConfigYml.getVikingDbHost(), tosProperties.getTosRegionName(), tosProperties.getTosSecretId(), tosProperties.getTosSecretKey(), "https");
        indexClient = new IndexClient(volcConfigYml.getCollectionName(), volcConfigYml.getIndexName(), volcConfigYml.getVikingDbHost(), tosProperties.getTosRegionName(), tosProperties.getTosSecretId(), tosProperties.getTosSecretKey(), "https");
    }

    public Map<String, Object> embeddingV2(String tosUrl) {
        try {
            HashMap<String, Object> params = new HashMap<>();
            params.put("return_token_usage", true);
            ArrayList<RawData> rawDatas = new ArrayList<>();
            rawDatas.add(new RawData().setDataType("image").setImage(tosUrl).build());
            Map<String, Object> res = vikingDBService.embeddingV2(new EmbModel().setModelName("doubao-embedding-vision").setParams(params).build(), rawDatas);
            return res;
        } catch (Exception e) {
            log.error(e.getMessage());
            return null;
        }
    }

    public Map<String, Object> queryTextVector(String text) {
        try {
            HashMap<String, Object> params = new HashMap<>();
            params.put("return_token_usage", true);
            ArrayList<RawData> rawDatas = new ArrayList<>();
            rawDatas.add(new RawData().setDataType("text").setText(text).build());
            Map<String, Object> res = vikingDBService.embeddingV2(new EmbModel().setModelName("doubao-embedding-large-and-m3").setParams(params).build(), rawDatas);
            return res;
        } catch (Exception e) {
            log.error(e.getMessage());
            return null;
        }
    }


    public void upsertData(List<VikingDbUserFileEntity> entityList) {
        if (CollectionUtils.isEmpty(entityList)) {
            return;
        }
        int size = entityList.size();
        int limit = 10;
        for (int i = 0; i < entityList.size(); i += limit) {
            List<VikingDbUserFileEntity> sub = entityList.subList(i, Math.min(i + limit, size));
            List<DataObject> dataObjectList = VikingDbConvert.convertDataObjectList(sub);
            doUpsertData(dataObjectList);
        }
    }

    public void upsertData(VikingDbUserFileEntity entity) {
        DataObject dataObject = VikingDbConvert.convertDataObject(entity);
        doUpsertData(dataObject);
    }


    public List<UserFileVO> searchWithMultiModal(GetSearchDataModel dataModel) {
        if (CollectionUtils.isEmpty(dataModel.getTosLinkList())) {
            return doSearchWithMultiModal(dataModel, null);
        }
        List<UserFileVO> res = new ArrayList<>();
        for (String tosLink : dataModel.getTosLinkList()) {
            List<UserFileVO> userFileVOList = doSearchWithMultiModal(dataModel, tosLink);
            if (CollectionUtils.isNotEmpty(userFileVOList)) {
                res.addAll(userFileVOList);
            }
        }
        return res;
    }

    private List<UserFileVO> doSearchWithMultiModal(GetSearchDataModel dataModel, String tosLink) {
        SearchWithMultiModalParam param = VikingDbConvert.getSearchWithMultiModalParam(dataModel, tosLink);
        try {
            log.info("多模态查询参数: {}", JSON.toJSONString(param));
            List<DataObject> dataObjectList = indexClient.searchWithMultiModal(param);
            List<UserFileVO> res = VikingDbConvert.convertUserFileVOList(dataObjectList);
            log.info("doSearchWithMultiModal: {}", JsonUtil.obj2Json(dataObjectList));
            return res;
        } catch (Exception e) {
            log.error("searchWithMultiModal", e);
            return null;
        }
    }


    public List<UserFileVO> search(GetSearchDataModel dataModel) {
        SearchParam param = VikingDbConvert.getSearchParam(dataModel);
        try {
            List<DataObject> dataObjectList = indexClient.search(param);
            return VikingDbConvert.convertUserFileVOList(dataObjectList);
        } catch (Exception e) {
            log.error("search", e);
            return null;
        }
    }

    public void delete(String userId) throws Exception {
        GetSearchDataModel dataModel = new GetSearchDataModel();
        dataModel.setUserId(userId);
        dataModel.setLimit(100);
        SearchParam param = VikingDbConvert.getSearchParam(dataModel);
        param.setOutputFields(null);
        List<DataObject> dataObjects = indexClient.search(param);
        for (DataObject dataObject : dataObjects) {
            collectionClient.deleteData(dataObject.getFields().get("__AUTO_ID__"));
        }
    }


    private void doUpsertData(List<DataObject> dataObjectList) {
        try {
            String bodyStr = JSON.toJSONString(dataObjectList);
            final String key = SecureUtil.md5(bodyStr);
            log.info("批量 doUpsertData: {}", bodyStr);
            long start = System.currentTimeMillis();
            RetryUtil.runWithThrow(key, () -> upsertDataThrowException(dataObjectList), 5, 4000);
            long end = System.currentTimeMillis();
            log.info("批量 doUpsertData cost: {}秒", (end - start) / 1000.0f);
        } catch (Exception e) {
            log.error("批量 doUpsertData", e);
        }
    }


    private void doUpsertData(DataObject dataObject) {
        try {
            String bodyStr = JSON.toJSONString(dataObject);
            final String key = SecureUtil.md5(bodyStr);
            log.info("单条 doUpsertData: {}", bodyStr);
            long start = System.currentTimeMillis();
            RetryUtil.runWithThrow(key, () -> upsertDataThrowException(dataObject), 5, 4000);
            long end = System.currentTimeMillis();
            log.info("单条 doUpsertData cost: {}秒", (end - start) / 1000.0f);
        } catch (Exception e) {
            log.error("单条 doUpsertData", e);
        }
    }

    private void upsertDataThrowException(List<DataObject> dataObjectList) {
        try {
            collectionClient.upsertData(dataObjectList);
        } catch (Exception e) {
            throw new RuntimeException(e);
        }
    }

    private void upsertDataThrowException(DataObject dataObject) {
        try {
            collectionClient.upsertData(dataObject);
        } catch (Exception e) {
            throw new RuntimeException(e);
        }
    }

}
